#include "inputimage.h"
#include "imagefeature.h"
#include <boost/foreach.hpp>
#include <boost/lexical_cast.hpp>
#include "commandlineoptions.h"
#include "stitcher.h"
#include "sift/utils.h"

using namespace std;
using namespace cshotel;
using namespace boost;
using namespace cv;
using namespace sift;

InputImage::InputImage(const std::string& name, const Size2i& size, const std::string& matchingVersionName, double scaleFactor, int newId) : fileName(name), imageSize(size), matchingVersionFileName(matchingVersionName), matchingVersionScaleFactor(scaleFactor), id(newId) { }

InputImage::~InputImage() {
}

const string& InputImage::getName() const {
    return fileName;
}

const string& InputImage::getMatchingVersionName() const {
    return matchingVersionFileName;
}

shared_ptr<Mat> InputImage::getImage() {
    if(shared_ptr<Mat> img = image.lock()) {
        return img;
    } else {
        shared_ptr<Mat> newImage(new Mat());
        Mat loadedImage = imread(fileName.c_str());
        *newImage = loadedImage;
        image = newImage;
        return newImage;
    }
}

shared_ptr<Mat> InputImage::getMatchingVersionImage() {
    if(shared_ptr<Mat> img = matchingVersionImage.lock()) {
        return img;
    } else {
        shared_ptr<Mat> newImage(new Mat());
        Mat loadedImage = imread(matchingVersionFileName.c_str());
        *newImage = loadedImage;
        matchingVersionImage = newImage;
        return newImage;
    }
}

int InputImage::getId() const {
    return id;
}

const Size2i& InputImage::getImageSize() {
    return imageSize;
}

const std::vector<shared_ptr<ImageFeature> >& InputImage::getImageFeatures() const {
    return features;
}

void InputImage::addImageFeature(shared_ptr<ImageFeature> newFeature) {
    features.push_back(newFeature);
}

void InputImage::addImageMatch(shared_ptr<ImageMatch> newMatch) {
    matches.push_back(newMatch);
}

const vector<shared_ptr<ImageMatch> >& InputImage::getImageMatches() const {
    return matches;
}

int InputImage::getImageId(weak_ptr<InputImage> img) {
    int id = -1;
    if(shared_ptr<InputImage> p = (img.lock())) {
        id = p->getId();
    }
    return id;
}

Mat InputImage::getGlobalHomography() const {
    Mat scale(Mat::eye(3, 3, CV_64FC1));
    scale.at<double>(0, 0) = scale.at<double>(1, 1) = matchingVersionScaleFactor;
    Mat rescale(Mat::eye(3, 3, CV_64FC1));
    rescale.at<double>(0, 0) = rescale.at<double>(1, 1) = 1.0 / matchingVersionScaleFactor;
    return rescale * matchingVersionHomography * scale;
}

Mat InputImage::getMatchingVersionHomography() const {
    return matchingVersionHomography;
}

void InputImage::setMatchingVersionHomography(const Mat& homography) {
    matchingVersionHomography = homography;
    calculateTransformedSizeAndOffset();
}

void InputImage::setTransformedVersion(const std::string& fileName) {
    transformedVersionFileName = fileName;
}

const string& InputImage::getTransformedVersionFileName() {
    return transformedVersionFileName;
}

const Point2d& InputImage::getTransformedVersionOffset() {
    return transformedVersionOffset;
}

const Size2d& InputImage::getTransformedVersionImageSize() {
    return transformedVersionImageSize;
}

void InputImage::calculateTransformedSizeAndOffset() {
    // calculate size and position of the transformed image
    // from the input image size and the homography matrix
    Mat currentTransformation = getGlobalHomography();
    Mat upperLeft = (Mat_<double>(3,1) << 0, 0, 1);
    Mat upperRight = (Mat_<double>(3,1) << imageSize.width, 0, 1);
    Mat lowerLeft = (Mat_<double>(3,1) << 0, imageSize.height, 1);
    Mat lowerRight = (Mat_<double>(3,1) << imageSize.width, imageSize.height, 1);
    
    Mat transformedUpperLeft = currentTransformation * upperLeft;
    transformedUpperLeft /= transformedUpperLeft.at<double>(2, 0);
    Mat transformedUpperRight = currentTransformation * upperRight;
    transformedUpperRight /= transformedUpperRight.at<double>(2, 0);
    Mat transformedLowerLeft = currentTransformation * lowerLeft;
    transformedLowerLeft /= transformedLowerLeft.at<double>(2, 0);
    Mat transformedLowerRight = currentTransformation * lowerRight;
    transformedLowerRight /= transformedLowerRight.at<double>(2, 0);
    
    double minX = min(transformedUpperLeft.at<double>(0, 0), transformedLowerLeft.at<double>(0, 0));
    double maxX = max(transformedUpperRight.at<double>(0, 0), transformedLowerRight.at<double>(0, 0));
    double minY = min(transformedUpperLeft.at<double>(1, 0), transformedUpperRight.at<double>(1, 0));
    double maxY = max(transformedLowerLeft.at<double>(1, 0), transformedLowerRight.at<double>(1, 0));
    
    transformedVersionOffset.x = minX;
    transformedVersionOffset.y = minY;
    transformedVersionImageSize.width = maxX - minX;
    transformedVersionImageSize.height = maxY - minY;
}

void InputImage::writeDebugImages(weak_ptr<InputImage> img) {
    if(shared_ptr<InputImage> image = img.lock()) {
        map<weak_ptr<InputImage>, shared_ptr<map<weak_ptr<ImageFeature>, weak_ptr<ImageFeature> > > > matches;
        
        
        foreach(shared_ptr<ImageFeature> feature, image->getImageFeatures()) {
            reverse_foreach(weak_ptr<ImageFeature> match, feature->getFeatureMatches()) {
                if(shared_ptr<ImageFeature> featureMatch = match.lock()) {
                    if(matches.find(featureMatch->getImage()) ==  matches.end()) {
                        shared_ptr<map<weak_ptr<ImageFeature>, weak_ptr<ImageFeature> > > newMap(new map<weak_ptr<ImageFeature>, weak_ptr<ImageFeature> >());
                        matches[featureMatch->getImage()] = newMap;
                    }
                    
                    (*(matches[featureMatch->getImage()]))[feature] = match;
                }
            }
        }
        
        typedef pair<weak_ptr<InputImage>, shared_ptr<map<weak_ptr<ImageFeature>, weak_ptr<ImageFeature> > > > ImageMatchPair;
        typedef pair<weak_ptr<ImageFeature>, weak_ptr<ImageFeature> > FeaturePair;
        foreach(ImageMatchPair match, matches) {
            vector<FeaturePair> matches;
            foreach(FeaturePair pair, *(match.second)) {
                matches.push_back(pair);
            }
            drawFeatureMatchImage(img, match.first, matches, "matches_");
        }
        
        
        shared_ptr<Mat> imageFile = image->getMatchingVersionImage();
        IplImage iplImage = static_cast<IplImage>(*imageFile);
        IplImage* featureImage = cvCloneImage(&iplImage);
        
        foreach(shared_ptr<ImageFeature> feature, image->getImageFeatures()) {
            CvScalar color = CV_RGB(rand() % 256, rand() % 256, rand() % 256);
            
            CvPoint p = cvPoint(feature->getSIFTFeature().x, feature->getSIFTFeature().y);
            cvCircle(featureImage, p, 2, color, -2);
            
            /*CvPoint p1 = cvPoint(feature->getSIFTFeature().x - 1, feature->getSIFTFeature().y - 1);
            CvPoint p2 = cvPoint(feature->getSIFTFeature().x + 1, feature->getSIFTFeature().y + 1);
            CvPoint p3 = cvPoint(feature->getSIFTFeature().x + 1, feature->getSIFTFeature().y - 1);
            CvPoint p4 = cvPoint(feature->getSIFTFeature().x - 1, feature->getSIFTFeature().y + 1);
            cvLine(featureImage, p1, p2, color);
            cvLine(featureImage, p3, p4, color);*/
        }
        
        string fileName = CommandLineOptions::getInstance().getWorkingDirectory() + "features_" + lexical_cast<string>(image->getId()) + ".jpg";
        cvSaveImage(fileName.c_str(), featureImage);
        cvReleaseImage(&featureImage);
    }
}

void InputImage::drawFeatureMatchImage(weak_ptr<InputImage> image1, weak_ptr<InputImage> image2, const vector<pair<weak_ptr<ImageFeature>, weak_ptr<ImageFeature> > >& features, string prefix) {
    if(shared_ptr<InputImage> img1 = image1.lock()) {
        if(shared_ptr<InputImage> img2 = image2.lock()) {
            shared_ptr<Mat> firstImageFile = img1->getMatchingVersionImage();
            shared_ptr<Mat> secondImageFile = img2->getMatchingVersionImage();
            IplImage firstIplImage = static_cast<IplImage>(*firstImageFile);
            IplImage secondIplImage = static_cast<IplImage>(*secondImageFile);
            IplImage* stackedImage = stack_imgs(&firstIplImage, &secondIplImage);
            
            typedef pair<weak_ptr<ImageFeature>, weak_ptr<ImageFeature> > FeaturePair;
            foreach(FeaturePair featurePair, features) {
                if(shared_ptr<ImageFeature> firstFeature = featurePair.first.lock()) {
                    if(shared_ptr<ImageFeature> secondFeature = featurePair.second.lock()) {
                        CvPoint p1 = cvPoint(firstFeature->getSIFTFeature().x, firstFeature->getSIFTFeature().y);
                        CvPoint p2 = cvPoint(secondFeature->getSIFTFeature().x, secondFeature->getSIFTFeature().y + firstIplImage.height);
                        cvLine(stackedImage, p1, p2, CV_RGB(rand() % 256, rand() % 256, rand() % 256));
                    }
                    
                }
            }
            
            string fileName = CommandLineOptions::getInstance().getWorkingDirectory() + prefix + lexical_cast<string>(img1->getId()) + "_" + lexical_cast<string>(img2->getId()) + ".jpg";
            cvSaveImage(fileName.c_str(), stackedImage);
            cvReleaseImage(&stackedImage);
        }
    }
}
